#ifndef TEST_MATRIX_UTIL
#define TEST_MATRIX_UTIL

#include "log_helper.hpp"
#include <new>

#define PRINTF_ARRAY_AS_MATRIX(os, mat, row, col)                              \
  do {                                                                         \
    for (auto i = 0; i < row; i++) {                                           \
      for (auto j = 0; j < col; j++) {                                         \
        os << mat[i * col + j] << " ";                                         \
      }                                                                        \
      os << std::endl;                                                         \
    }                                                                          \
    os << std::endl;                                                           \
  } while (0)

#define ALLOC_INT_ARRAY_WITHOUT_INIT(arr, len)                                 \
  do {                                                                         \
    try {                                                                      \
      arr = new int[len];                                                      \
    } catch (const std::bad_array_new_length &e) {                             \
      ERROR_PRINTF("%s in ALLOC_INT_ARRAY_WITHOUT_INIT with len=%zu\n",        \
                   e.what(), len);                                             \
    }                                                                          \
  } while (0)

#define ALLOC_INT_ARRAY_WITH_INIT(arr, len)                                    \
  do {                                                                         \
    try {                                                                      \
      arr = new int[len]();                                                    \
    } catch (const std::bad_array_new_length &e) {                             \
      ERROR_PRINTF("%s in ALLOC_INT_ARRAY_WITH_INIT with len=%zu\n", e.what(), \
                   len);                                                       \
    }                                                                          \
  } while (0)

#define ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(arr, len)                              \
  do {                                                                         \
    try {                                                                      \
      arr = new double[len];                                                   \
    } catch (const std::bad_array_new_length &e) {                             \
      ERROR_PRINTF("%s in ALLOC_DOUBLE_ARRAY_WITHOUT_INIT with len=%zu\n",     \
                   e.what(), len);                                             \
    }                                                                          \
  } while (0)

#define ALLOC_DOUBLE_ARRAY_WITH_INIT(arr, len)                                 \
  do {                                                                         \
    try {                                                                      \
      arr = new double[len]();                                                 \
    } catch (const std::bad_array_new_length &e) {                             \
      ERROR_PRINTF("%s in ALLOC_DOUBLE_ARRAY_WITH_INIT with len=%zu\n",        \
                   e.what(), len);                                             \
    }                                                                          \
  } while (0)

#define COPY_ARRAY_I2D(in, out, len)                                           \
  do {                                                                         \
    for (auto i = 0; i < len; i++) {                                           \
      (out)[i] = static_cast<double>((in)[i]);                                 \
    }                                                                          \
  } while (0)

#define COPY_ARRAY_D2I(in, out, len)                                           \
  do {                                                                         \
    for (auto i = 0; i < len; i++) {                                           \
      (out)[i] = static_cast<int>((in)[i]);                                    \
    }                                                                          \
  } while (0)

#define COPY_ARRAY_TO_NTL_MATRIX(ntl, arr)                                     \
  do {                                                                         \
    for (auto i = 0; i < ntl.NumRows(); i++) {                                 \
      for (auto j = 0; j < ntl.NumCols(); j++) {                               \
        ntl[i][j] = arr[i * ntl.NumCols() + j];                                \
      }                                                                        \
    }                                                                          \
  } while (0)

#define SCALE_DOUBLE_ARRAY(arr, len, scale)                                    \
  do {                                                                         \
    for (int i = 0; i < (len); i++) {                                          \
      (arr)[i] *= (scale);                                                     \
    }                                                                          \
  } while (0)

//* true or false
#define FIX_RAND_SEED(enable)                                                  \
  do {                                                                         \
    std::srand(enable ? 0l : std::time(nullptr));                              \
  } while (0)

//! Need allocate memory mannully
//! bl <= 32
#define RAND_ARRAY(arr, len, bl)                                               \
  do {                                                                         \
    unsigned int mask = (1u << bl) - 1;                                        \
    for (int i = 0; i < (len); i++) {                                          \
      (arr)[i] = std::rand() & mask;                                           \
    }                                                                          \
  } while (0)

//* true or false
#define FIX_NTL_RAND_SEED(enable)                                              \
  do {                                                                         \
    NTL::SetSeed(NTL::conv<NTL::ZZ>(enable ? 0l : std::time(nullptr)));        \
  } while (0)

#define RAND_NTL_MATRIX(mat, row, col, bl)                                     \
  do {                                                                         \
    (mat).SetDims(row, col);                                                   \
    for (long i = 0; i < row; i++) {                                           \
      for (long j = 0; j < col; j++) {                                         \
        (mat)[i][j] = NTL::RandomBits_ZZ(bl);                                  \
      }                                                                        \
    }                                                                          \
  } while (0)

#ifdef USE_INTEL_MKL
#include "mkl.h"

//! ||C||_\infty <= 2^{31}
static void cblas_igemm(const CBLAS_LAYOUT Layout, const CBLAS_TRANSPOSE TransA,
                        const CBLAS_TRANSPOSE TransB, const long long M,
                        const long long N, const long long K, const int alpha,
                        const int *A, const long long lda, const int *B,
                        const long long ldb, const int beta, int *C,
                        const long long ldc) {
  const std::size_t size_A = M * K;
  const std::size_t size_B = K * N;
  const std::size_t size_C = M * N;

  double *DA = nullptr;
  double *DB = nullptr;
  double *DC = nullptr;
  ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DA, size_A);
  ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DB, size_B);
  ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DC, size_C);
  COPY_ARRAY_I2D(A, DA, size_A);
  COPY_ARRAY_I2D(B, DB, size_B);
  cblas_dgemm(Layout, TransA, TransB, M, N, K, static_cast<double>(alpha), DA,
              lda, DB, ldb, static_cast<double>(beta), DC, ldc);
  COPY_ARRAY_D2I(DC, C, size_C);

  delete[] DA;
  delete[] DB;
  delete[] DC;
}

#endif
#endif /* TEST_MATRIX_UTIL */
